#!/usr/bin/env python3
from typing import Tuple, List, Optional, Union  # 用于类型注解的工具
# paramiko 是一个用于实现 SSH2 协议的 Python 库
from paramiko import SSHClient as ParamikoSSHClient, AutoAddPolicy, RSAKey
from io import StringIO  # 用于在内存中处理字符串
from subprocess import Popen, PIPE  # 用于执行本地系统命令
from datetime import datetime  # 用于处理日期和时间
from pathlib import Path  # 用于处理文件路径


class SSHConfig:
    """SSH 连接配置类
    
    每个属性都使用类型注解来指定其类型，这样可以：
    1. 提供更好的代码提示
    2. 帮助开发工具进行类型检查
    3. 使代码更容易理解和维护
    
    属性说明：
        hostname: 要连接的远程主机地址（IP 或域名）
        port: SSH 服务端口号，SSH 默认是 22 端口
        username: 登录用户名，默认是 root 用户
        password: SSH 登录密码，如果使用密钥登录则不需要
        pkey: SSH 私钥内容，可以直接传入密钥字符串
        private_key_file: SSH 私钥文件路径，可以指定密钥文件位置
        pkey_pwd: SSH 私钥的密码（如果私钥设置了密码保护）
        sock: 跳板机的 socket 对象（用于通过跳板机连接目标服务器）
    """
    def __init__(self, hostname: str, port: int = 22, username: str = 'root',
                 password: Optional[str] = None, pkey: Optional[str] = None,
                 private_key_file: Optional[Union[str, Path]] = None,
                 pkey_pwd: Optional[str] = None, sock: Optional[object] = None):
        self.hostname = hostname  # 主机地址，必填项
        self.port = port  # 端口号，默认 22
        self.username = username  # 用户名，默认 root
        self.password = password  # 密码，可选
        self.pkey = pkey  # 密钥内容，可选
        self.private_key_file = private_key_file  # 密钥文件路径，可选
        self.pkey_pwd = pkey_pwd  # 密钥密码，可选
        self.sock = sock  # 跳板机 socket，可选

    def __repr__(self) -> str:
        return f"SSHConfig(hostname='{self.hostname}', port={self.port}, username='{self.username}')"


class Client:
    """SSH 客户端类
    
    这个类封装了 paramiko 的 SSHClient，提供了更简单的接口来：
    1. 建立 SSH 连接
    2. 执行远程命令
    3. 管理 SFTP 会话
    4. 处理连接的开启和关闭
    """
    
    def __init__(self, **kwargs):
        """初始化 SSH 客户端
        
        参数:
            **kwargs: 关键字参数，将被传递给 SSHConfig 来创建配置对象
                     支持的参数见 SSHConfig 类的属性说明
        
        初始化过程：
        1. 创建配置对象
        2. 初始化 SSH 和 SFTP 连接为 None
        3. 立即尝试建立连接
        """
        self.config = SSHConfig(**kwargs)  # 使用传入的参数创建配置对象
        self.ssh = None  # SSH 连接对象
        self.sftp = None  # SFTP 会话对象
        self._connect()  # 建立连接

    def _get_pkey(self) -> Optional[RSAKey]:
        """获取 SSH 私钥对象
        
        这个方法会尝试从两个来源获取私钥：
        1. 直接提供的私钥内容（self.config.pkey）
        2. 私钥文件路径（self.config.private_key_file）
        
        返回:
            RSAKey: paramiko 的 RSA 密钥对象
            None: 如果没有提供任何私钥
        """
        # 如果提供了私钥内容，从字符串创建密钥对象
        if self.config.pkey:
            return RSAKey.from_private_key(
                file_obj=StringIO(self.config.pkey),  # 将字符串转换为文件对象
                password=self.config.pkey_pwd  # 私钥密码（如果有）
            )
        # 如果提供了私钥文件路径，从文件创建密钥对象
        if self.config.private_key_file:
            key_path = Path(self.config.private_key_file)
            if key_path.exists():  # 确保文件存在
                return RSAKey.from_private_key_file(
                    str(key_path),
                    password=self.config.pkey_pwd
                )
        return None

    def _connect(self) -> None:
        """建立 SSH 连接
        
        这个方法会：
        1. 验证必要的配置参数
        2. 创建 SSH 客户端并设置密钥策略
        3. 建立连接并创建 SFTP 会话
        4. 处理连接超时和其他异常
        
        如果连接失败，会关闭已建立的连接并抛出异常
        """
        # 检查必要的配置参数
        if not all([self.config.hostname, self.config.username]):
            raise Exception("hostname 和 username 不能为空")

        try:
            # 获取私钥对象
            pkey = self._get_pkey()
            
            # 创建 SSH 客户端对象
            self.ssh = ParamikoSSHClient()
            # 设置密钥策略：自动添加主机密钥
            self.ssh.set_missing_host_key_policy(AutoAddPolicy())

            # 记录开始时间，用于检查是否超时
            start_time = datetime.now()
            
            # 建立 SSH 连接
            self.ssh.connect(
                hostname=self.config.hostname.strip(),  # 移除主机名两端的空白
                port=self.config.port,
                username=self.config.username.strip(),
                password=self.config.password,
                pkey=pkey,  # RSA 密钥对象
                sock=self.config.sock,  # 跳板机 socket
                timeout=5,  # 连接超时时间
                auth_timeout=5,  # 认证超时时间
                allow_agent=False,  # 不使用 SSH 代理
                look_for_keys=False  # 不在默认位置查找密钥
            )
            
            # 创建 SFTP 会话
            self.sftp = self.ssh.open_sftp()

            # 检查是否连接超时（超过 60 秒）
            connect_time = (datetime.now() - start_time).seconds
            if connect_time > 60:
                raise Exception("主机连接超时，请检查网络状况或主机状态")

        except Exception as e:
            # 如果发生异常，确保关闭所有连接
            self.close()
            raise Exception(f"连接失败 {self.config.hostname}: {str(e)}")

    def remote_command(self, command: str, error: bool = False) -> Tuple[int, str, str]:
        """执行远程命令
        
        在远程服务器上执行 shell 命令，并获取执行结果。
        
        参数:
            command: 要执行的 shell 命令
            error: 当命令执行失败时是否抛出异常
        
        返回:
            tuple: (退出码, 标准输出, 标准错误)
                - 退出码: 0 表示成功，非 0 表示失败
                - 标准输出: 命令的输出内容
                - 标准错误: 命令的错误信息
        """
        # 执行远程命令
        stdin, stdout, stderr = self.ssh.exec_command(command, timeout=5)

        # 获取命令执行结果
        state = stdout.channel.recv_exit_status()  # 退出状态码
        stdout_str = stdout.read().decode()  # 标准输出
        stderr_str = stderr.read().decode()  # 标准错误

        # 如果设置了 error=True 且命令执行失败，抛出异常
        if error and state != 0:
            raise Exception(f"命令执行失败: {command}\n错误信息: {stderr_str}")

        return state, stdout_str, stderr_str

    def close(self) -> None:
        """关闭 SSH 连接
        
        按照正确的顺序关闭 SFTP 会话和 SSH 连接。
        这个方法应该在完成所有操作后调用，以释放资源。
        """
        if self.sftp:
            self.sftp.close()  # 先关闭 SFTP 会话
        if self.ssh:
            self.ssh.close()  # 再关闭 SSH 连接


class JumpServer:
    """跳板机客户端类
    
    用于通过跳板机连接目标服务器。跳板机是一个中间服务器，
    用户先连接到跳板机，然后通过跳板机连接到目标服务器。
    这种方式通常用于访问不能直接连接的内网服务器。
    """
    
    def __init__(self, dest_addr: List, **kwargs):
        """初始化跳板机客户端
        
        参数:
            dest_addr: 目标服务器的地址和端口，格式: [host, port]
            **kwargs: 跳板机的连接参数，与 SSHConfig 的参数相同
        """
        self.config = SSHConfig(**kwargs)  # 跳板机的配置
        self.dest_addr = dest_addr  # 目标服务器地址
        self.ssh = None  # SSH 连接对象
        self.channel = None  # 通道对象
        self._connect()  # 建立连接

    def _connect(self) -> None:
        """建立跳板机连接
        
        这个方法会：
        1. 连接到跳板机
        2. 创建一个通道连接到目标服务器
        
        如果连接失败，会清理资源并抛出异常
        """
        if not all([self.config.hostname, self.config.username]):
            raise Exception("hostname 和 username 不能为空")

        try:
            # 获取 SSH 私钥
            pkey = None
            if self.config.pkey:
                pkey = RSAKey.from_private_key(
                    file_obj=StringIO(self.config.pkey),
                    password=self.config.pkey_pwd
                )
            elif self.config.private_key_file:
                key_path = Path(self.config.private_key_file)
                if key_path.exists():
                    pkey = RSAKey.from_private_key_file(
                        str(key_path),
                        password=self.config.pkey_pwd
                    )

            # 创建并配置 SSH 客户端
            self.ssh = ParamikoSSHClient()
            self.ssh.set_missing_host_key_policy(AutoAddPolicy())

            # 连接到跳板机
            self.ssh.connect(
                hostname=self.config.hostname.strip(),
                port=self.config.port,
                username=self.config.username.strip(),
                password=self.config.password,
                pkey=pkey,
                timeout=5,
                auth_timeout=5,
                allow_agent=False,
                look_for_keys=False
            )

            # 获取传输对象并设置超时
            transport = self.ssh.get_transport()
            transport.banner_timeout = 30
            
            # 创建通道连接到目标服务器
            self.channel = transport.open_channel(
                kind="direct-tcpip",  # TCP 直接通道
                dest_addr=self.dest_addr,  # 目标服务器地址和端口
                src_addr=(self.config.hostname, self.config.port)  # 源地址（跳板机）
            )

        except Exception as e:
            self.close()
            raise Exception(f"跳板机连接失败: {str(e)}")

    def close(self) -> None:
        """关闭跳板机连接
        
        关闭 SSH 连接并释放资源
        """
        if self.ssh:
            self.ssh.close()


def local_command(command: str, error: bool = False) -> Tuple[int, str, str]:
    """执行本地系统命令
    
    使用 subprocess.Popen 在本地执行 shell 命令。
    
    参数:
        command: 要执行的 shell 命令
        error: 命令执行失败时是否抛出异常
    
    返回:
        tuple: (退出码, 标准输出, 标准错误)
            - 退出码: 0 表示成功，非 0 表示失败
            - 标准输出: 命令的输出内容
            - 标准错误: 命令的错误信息
    """
    # 创建子进程执行命令
    process = Popen(
        command,
        stdout=PIPE,  # 捕获标准输出
        stderr=PIPE,  # 捕获标准错误
        shell=True  # 使用 shell 执行命令
    )

    # 等待命令执行完成并获取输出
    stdout, stderr = process.communicate()
    state = process.returncode  # 获取退出码
    
    # 解码输出内容
    stderr_str = stderr.decode('utf-8', 'ignore')  # 忽略无法解码的字符
    stdout_str = stdout.decode('utf-8')

    # 如果设置了 error=True 且命令执行失败，抛出异常
    if error and state != 0:
        raise Exception(f"本地命令执行失败: {command}\n错误信息: {stderr_str}")

    return state, stdout_str, stderr_str


# 如果直接运行此文件，执行示例代码
if __name__ == "__main__":
    # 创建跳板机连接
    jumpServer = JumpServer(
        hostname="192.168.110.28",  # 跳板机地址
        port=5423,  # 跳板机端口
        private_key_file="/home/ubuntu/.ssh/id_rsa",  # 私钥文件
        username="ubuntu",  # 用户名
        dest_addr=["192.168.110.32", 5423]  # 目标服务器地址和端口
    )
    channel = jumpServer.channel  # 获取通道

    # 通过跳板机连接目标服务器
    client = Client(
        hostname="192.168.110.32",  # 目标服务器地址
        port=5423,  # 目标服务器端口
        private_key_file="/home/ubuntu/.ssh/id_rsa",  # 私钥文件
        username="ubuntu",  # 用户名
        sock=channel  # 使用跳板机通道
    )

    # 在目标服务器上执行 ls 命令
    state, stdout_str, stderr_str = client.remote_command("ifconfig")
    if state == 0:
        print(stdout_str)  # 打印命令输出
    else:
        print(stderr_str)  # 打印错误信息

    # 清理资源
    if client:
        client.close()  # 关闭目标服务器连接

    if jumpServer:
        jumpServer.close()  # 关闭跳板机连接